/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package com.xiashitech.agent.instrumentation.redis.redisson;

import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;

import java.util.concurrent.CompletableFuture;

public final class XSCompletableFutureWrapper<T> extends XSCompletableFuture<T>{

    private XSCompletableFutureWrapper(CompletableFuture<T> delegate) {
        Context context = Context.current();
        this.whenComplete(
                (result, error) -> {
                    XSEndOperationListener<T> endOperationListener = getEndOperationListener();
                    if (endOperationListener != null) {
                        redisData = endOperationListener.apply(result, error);
                    }
                    Context delegateContext = BaggageUtil.syncBaggage().build().storeInContext(context);
                    try (Scope ignored = delegateContext.makeCurrent()) {
                        if (error != null) {
                            delegate.completeExceptionally(error);
                        } else {
                            delegate.complete(result);
                        }
                    }
                });
    }

    public static <T> CompletableFuture<T> wrap(CompletableFuture<T> delegate) {
        if (delegate instanceof XSCompletableFuture) {
            return delegate;
        }
        return new XSCompletableFutureWrapper(delegate);
    }
}
